# 导入函数库
from jqdata import *


# 初始化函数，设定基准等等
def initialize(context):
    # 设定沪深300作为基准
    set_benchmark('000300.XSHG')
    # 开启动态复权模式(真实价格)
    set_option('use_real_price', True)
    # 输出内容到日志 log.info()
    log.info('初始函数开始运行且全局只运行一次')
    # 过滤掉order系列API产生的比error级别低的log
    log.set_level('order', 'error')

    ### 股票相关设定 ###
    # 股票类每笔交易时的手续费是：买入时佣金万分之三，卖出时佣金万分之三加千分之一印花税, 每笔交易佣金最低扣5块钱
    set_order_cost(OrderCost(close_tax=0.001, open_commission=0.00015, close_commission=0.00015, min_commission=5),
                   type='stock')

    # 开盘时运行
    run_daily(sell_trade, time='09:45')

    # 开盘时运行
    run_daily(trade, time='14:51')


# 卖出策略
def sell_trade(context):
    long_positions_dict = context.portfolio.long_positions
    for position in list(long_positions_dict.values()):
        # 取前6天的数据
        bars_array_7 = get_bars(position.security, count=11, unit='1d', include_now=True, fields=['close'])['close']
        # 取最近5天的平均价
        pre_ma5 = bars_array_7[0:10].mean()
        ma5 = bars_array_7[-10:].mean()
        # 今天5日线比昨天低，卖
        if ma5 < pre_ma5:
            order_target(position.security, 0)
            log.info("09:45 %s 5日线比昨天低，ma5：%f，pre_ma5：%f=============卖出" % (position.security, ma5, pre_ma5))


# 策略
def trade(context):
    # 看看沪深300趋势
    hs300_bars_array = get_bars('000300.XSHG', count=11, unit='1d', include_now=True, fields=['close'])['close']
    hs300_pre_ma = hs300_bars_array[0:10].mean()
    hs300_ma = hs300_bars_array[-10:].mean()

    # 先卖出股票腾仓位
    # 看看持有股票，如果有5日线下穿20日线，卖出
    long_positions_dict = context.portfolio.long_positions
    for position in list(long_positions_dict.values()):
        # 取前6天的数据
        bars_array_7 = get_bars(position.security, count=11, unit='1d', include_now=True, fields=['close'])['close']
        # 取最近5天的平均价
        pre_ma5 = bars_array_7[0:10].mean()
        ma5 = bars_array_7[-10:].mean()
        # 今天5日线比昨天低，卖
        if ma5 < pre_ma5:
            order_target(position.security, 0)
            log.info("14:50 %s 5日线比昨天低，ma5：%f，pre_ma5：%f=============卖出" % (position.security, ma5, pre_ma5))

        # if hs300_ma < hs300_pre_ma:
        #     order_target(position.security, 0)
        #     log.info("14:50 沪深300均线比昨天低，hs300_ma：%f，hs300_pre_ma：%f=============卖出" % (hs300_ma, hs300_pre_ma))

    # if hs300_ma < hs300_pre_ma:
    #     log.info("14:50 沪深300均线比昨天低，hs300_ma：%f，hs300_pre_ma：%f=============今天不买了" % (hs300_ma, hs300_pre_ma))
    #     return

    # 如果没有仓位了，今天不动了
    available_cash = context.portfolio.available_cash
    if available_cash < 3000:
        log.info("没钱了，今天不动了")
        return

    # 账户总资产（现金+股票市值）, 每次买入一只票最多5成仓
    buy_max_per_stock = round(context.portfolio.total_value * 0.5, 2)
    # 前一天
    pre_date = context.current_dt + timedelta(days=-1)
    # 在备选股中寻找5日线上穿20日线的票
    stocks = top_money_stock(context)
    # stocks = top_main_money_stock(context)
    # stocks = get_smallest(context, stocks)
    log.info("今日备选===========" + str(stocks))
    for stock in stocks:
        # 取得当前的现金
        cash = context.portfolio.available_cash
        if cash < 3000:
            log.info("没钱了，今天不继续买了")
            return
        # 取过去6天的数据，因为还需要计算昨天的ma5
        bars_array_7 = get_bars(stock, count=11, unit='1d', include_now=True, fields=['close'])['close']
        pre_ma5 = bars_array_7[0:10].mean()
        ma5 = bars_array_7[-10:].mean()
        if ma5 > pre_ma5 and context.portfolio.positions[stock].closeable_amount <= 0:
            buy_cash = min(cash, buy_max_per_stock)
            order_value(stock, buy_cash)
            log.info("%s 5日线比昨天高，ma5：%f，pre_ma5：%f=============买入" % (stock, ma5, pre_ma5))


# 返回成交额前300的股票
def top_money_stock(context):
    # 获取所有上市股票
    all_stock = list(get_all_securities(['stock'], date=context.current_dt).index)
    # 获取所有股票当日成交额
    money = get_price(all_stock, context.current_dt, context.current_dt, fields=["money"], panel=False)
    # 排序取前200条
    sorted_money = money.sort_values(by='money', ascending=False).head(300)

    # 返回股票代码列表
    return list(sorted_money['code'])


# 获取主力净流入靠前的票
def top_main_money_stock(context):
    # 获取所有上市股票
    all_stock = list(get_all_securities(['stock'], date=context.current_dt).index)
    # 获取所有股票当日主力资金流向
    main_money_flow = get_money_flow(all_stock, context.current_dt, context.current_dt,
                                     ["date", "sec_code", "net_amount_main"])
    # 排序取前200条
    sorted_money_flow = main_money_flow.sort_values(by='net_amount_main', ascending=False).head(50)

    # 返回股票代码列表
    return list(sorted_money_flow['sec_code'])


# 选择市值最小的票
def get_smallest(context, stock_list):
    q = query(valuation.code, valuation.circulating_market_cap).filter(valuation.code.in_(stock_list)).order_by(
        valuation.circulating_market_cap.asc())
    df = get_fundamentals(q, date=context.current_dt)
    df.index = df.code
    stocks = list(df.code)[0:20]
    return stocks

